// Copyright (c) 2020 cgsdfc
//
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#pragma once

#include "pyast/pyast.h"

namespace pydsl {

class DslScope {
 public:
  static pyast::Arena* current();

  DslScope();
  ~DslScope();

 private:
  static pyast::Arena* g_arena_;
};

struct NoneType {};

extern const NoneType None;

inline pyast::expr_node CreateConstantImpl(pybind11::object value) {
  return pyast::ConstantNode{value, "", DslScope::current()};
}

inline pyast::expr_node CreateBoolConstant(bool value) {
  return CreateConstantImpl(pybind11::bool_{value});
}

inline pyast::expr_node CreateIntConstant(long value) {
  return CreateConstantImpl(pybind11::int_{value});
}

inline pyast::expr_node CreateFloatConstant(double value) {
  return CreateConstantImpl(pybind11::float_{value});
}

inline pyast::expr_node CreateNoneConstant(NoneType) {
  return CreateConstantImpl(pybind11::none{});
}

inline pyast::expr_node CreateStrConstant(const char* str) {
  return CreateConstantImpl(pybind11::str{str});
}

class Expression {
 public:
  Expression(pyast::expr_node node) : node_(node) {}

  Expression(bool value) : Expression{CreateBoolConstant(value)} {}
  Expression(double value) : Expression{CreateFloatConstant(value)} {}
  Expression(int value) : Expression{CreateIntConstant(value)} {}
  Expression(NoneType value) : Expression{CreateNoneConstant(value)} {}
  Expression(const char* value) : Expression{CreateStrConstant(value)} {}

  Expression operator()() const { return pyast::CallNode{}; }

  operator pyast::expr_node() const { return node_; }

 private:
  pyast::expr_node node_;
};

namespace literals {

inline Expression operator""_n(const char* name) {
  return {pyast::NameNode{pybind11::str{name}, pyast::expr_context_kind::kLoad,
                          DslScope::current()}};
}

inline Expression operator""_s(const char* string) {
  return Expression{string};
}

}  // namespace literals

#define PYDSL_GEN_BIN_OP(op, op_name)                                \
  inline Expression operator op(Expression lhs, Expression rhs) {    \
    return pyast::BinOpNode{lhs, pyast::operator_kind::op_name, rhs, \
                            DslScope::current()};                    \
  }

PYDSL_GEN_BIN_OP(+, kAdd)
PYDSL_GEN_BIN_OP(-, kSub)
PYDSL_GEN_BIN_OP(*, kMult)
PYDSL_GEN_BIN_OP(/, kDiv)
PYDSL_GEN_BIN_OP(%, kMod)
PYDSL_GEN_BIN_OP(<<, kLShift)
PYDSL_GEN_BIN_OP(>>, kRShift)
PYDSL_GEN_BIN_OP(|, kBitOr)
PYDSL_GEN_BIN_OP(^, kBitXor)
PYDSL_GEN_BIN_OP(&, kBitAnd)

#undef PYDSL_GEN_BIN_OP

class Statement {
 public:
  Statement(pyast::stmt_node node);
  Statement(Expression expr)
      : Statement{pyast::ExprNode{static_cast<pyast::expr_node>(expr),
                                  DslScope::current()}} {}

 private:
  pyast::stmt_node node_;
};

#define PYDSL_GEN_AUG_ASSIGN(op, op_name)                                     \
  inline Statement operator op(Expression target, Expression value) {         \
    return pyast::AugAssignNode{target, pyast::operator_kind::op_name, value, \
                                DslScope::current()};                         \
  }

PYDSL_GEN_AUG_ASSIGN(+=, kAdd)
PYDSL_GEN_AUG_ASSIGN(-=, kSub)
PYDSL_GEN_AUG_ASSIGN(*=, kMult)
PYDSL_GEN_AUG_ASSIGN(/=, kDiv)
PYDSL_GEN_AUG_ASSIGN(%=, kMod)
PYDSL_GEN_AUG_ASSIGN(<<=, kLShift)
PYDSL_GEN_AUG_ASSIGN(>>=, kRShift)
PYDSL_GEN_AUG_ASSIGN(|=, kBitOr)
PYDSL_GEN_AUG_ASSIGN(^=, kBitXor)
PYDSL_GEN_AUG_ASSIGN(&=, kBitAnd)

#undef PYDSL_GEN_AUG_ASSIGN

}  // namespace pydsl
